from math import log

'''
   决策树示例数据，二维数组:
   dataSet = [
        [1, 1, 'yes'],
        [1, 1, 'yes'],               
        [1, 0, 'no'],
        [0, 1, 'no'],
        [0, 1, 'no']
   ]
   元素含义：不浮出水面是否可以生存, 是否有脚蹼， 是否是鱼类
    
'''
def createDataSet():
    dataSet = [
       [1, 1, 'yes'],
       [1, 1, 'yes'],
       [1, 0, 'no'],
       [0, 1, 'no'],
       [0, 1, 'no']
    ]

    labels = ['no surfacing', 'flippers']
    return dataSet, labels

'''
    计算给定列表的香农熵
'''
def calcShannonEnt(dataSet):
    #拿到dataSet元素的个数
    numEntries = len(dataSet)
    labelCounts = {}

    #开始遍历
    for featVec in dataSet:
        #本元素数组最后一个值，即为label
        currentLabel = featVec[-1]
        #若label数组中无此label，初始化数量为0
        if currentLabel not in labelCounts.keys() : labelCounts[currentLabel] = 0
        #给对应label的数量加1
        labelCounts[currentLabel] += 1
    shannonEnt = 0.0

    #计算香农熵
    for key in labelCounts:
        prob = float(labelCounts[key])/numEntries
        shannonEnt -= prob * log(prob, 2)

    return shannonEnt

'''
    匹配dataSet中索引在axis处，值为value的元素，并将匹配后数组元素-剔除axis处位置的值
    若axis为1,value为1，则返回：
    [[1, 'yes'], [1, 'yes'], [0, 'no'], [0, 'no']]
    意为：在二维数组dataSet中找到索引在1处（索引从0开始），值为1的数组元素, 并将这些元素索引1处的元素去除并返回（新的二维数组）
'''
def splitDataSet(dataSet, axis, value):
    retDataSet = []
    for featVec in dataSet:
        if featVec[axis] == value:
            reducedFeatVec = featVec[:axis]
            reducedFeatVec.extend(featVec[axis+1:])
            retDataSet.append(reducedFeatVec)
    return retDataSet

'''
    选择最佳匹配特性
'''
def chooseBestFeatureToSplit(dataSet):
    #初始化特性的数量
    numFeatures = len(dataSet[0]) - 1

    #计算整体香农熵
    baseEntropy = calcShannonEnt(dataSet)

    #初始化信息增益值为0.0
    bestInfoGain = 0.0

    #最优分类特性索引-1
    bestFeature = -1

    for i in range(numFeatures):
        #只取索引i位置处的值，形成新的二维数组，例如:[[1][0][1][1]]
        featList = [example[i] for example in dataSet]
        #对i处特性值去重，得唯一特性值列表
        uniqueVals = set(featList)
        newEntropy = 0.0
        #计算索引i处对应特性的香农熵
        for value in uniqueVals:
            #查找匹配列表
            subDataSet = splitDataSet(dataSet, i, value)
            #计算value匹配列表占比
            prob = len(subDataSet)/float(len(dataSet))
            #计算i处特性的熵的和
            newEntropy += prob * calcShannonEnt(subDataSet)
        #计算信息增益
        infoGain = baseEntropy - newEntropy
        #找到信息增益最大值，即为最佳匹配，即不确定性最大
        if(infoGain > bestInfoGain):
            bestInfoGain = infoGain
            bestFeature = i

    return bestFeature


